import pygame
import random
import math

WIDTH, HEIGHT = 800, 600
BOIDS = 100

class Boid:
    def __init__(self, x, y):
        self.position = pygame.Vector2(x, y)
        self.velocity = pygame.Vector2(random.uniform(-1, 1), random.uniform(-1, 1)).normalize() * 3
        self.acceleration = pygame.Vector2(0, 0)

    def update(self, boids):
        self.apply_rules(boids)
        self.velocity += self.acceleration
        self.velocity.clamp_magnitude(3)
        self.position += self.velocity
        self.acceleration *= 0

    def apply_rules(self, boids):
        separation = self.separate(boids)
        alignment = self.align(boids)
        cohesion = self.cohesion(boids)

        self.acceleration += separation * 1.5
        self.acceleration += alignment * 1
        self.acceleration += cohesion * 1

    def separate(self, boids):
        steer = pygame.Vector2(0, 0)
        count = 0
        desired_separation = 25
        for boid in boids:
            distance = self.position.distance_to(boid.position)
            if distance > 0 and distance < desired_separation:
                diff = self.position - boid.position
                diff.normalize()
                diff /= distance
                steer += diff
                count += 1
        if count > 0:
            steer /= count
        if steer.magnitude() > 0:
            steer.normalize()
            steer *= 3
        return steer

    def align(self, boids):
        steer = pygame.Vector2(0, 0)
        count = 0
        neighbor_dist = 50
        for boid in boids:
            distance = self.position.distance_to(boid.position)
            if distance > 0 and distance < neighbor_dist:
                steer += boid.velocity
                count += 1
        if count > 0:
            steer /= count
        if steer.magnitude() > 0:
            steer.normalize()
            steer *= 3
        return steer

    def cohesion(self, boids):
        steer = pygame.Vector2(0, 0)
        count = 0
        neighbor_dist = 50
        for boid in boids:
            distance = self.position.distance_to(boid.position)
            if distance > 0 and distance < neighbor_dist:
                steer += boid.position
                count += 1
        if count > 0:
            steer /= count
        if steer.magnitude() > 0:
            steer = steer - self.position
            steer.normalize()
            steer *= 3
        return steer

    def show(self, screen):
        pygame.draw.circle(screen, (255, 255, 255), (int(self.position.x), int(self.position.y)), 3)

def main():
    pygame.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    clock = pygame.time.Clock()

    boids = [Boid(random.randrange(WIDTH), random.randrange(HEIGHT)) for _ in range(BOIDS)]

    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        screen.fill((0, 0, 0))

        for boid in boids:
            boid.update(boids)
            boid.show(screen)

        pygame.display.flip()
        clock.tick(60)

    pygame.quit()

if __name__ == "__main__":
    main()
